import numpy as np
import onnx
from onnx import shape_inference
import onnxmltools
from utils.general import check_requirements
try:
    check_requirements("nvidia-pyindex")
    check_requirements('onnx_graphsurgeon')
    import onnx_graphsurgeon as gs
except Exception as e:
    print('Import onnx_graphsurgeon failure: %s' % e)


class RegisterNMS(object):
    def __init__(self, logger, onnx_model_path: str, precision: str = "fp32", prefix=''):

        self.graph = gs.import_onnx(onnx.load(onnx_model_path))
        assert self.graph, 'grap must not empty'
        # Fold constants via ONNX-GS that PyTorch2ONNX may have missed
        self.graph.fold_constants()
        self.precision = precision
        self.batch_size = 1
        self.logger = logger
        self.prefix = prefix
        self.logger.info(f"{prefix} graph created successfully")

    def infer(self):
        """
        Sanitize the graph by cleaning any unconnected nodes, do a topological resort,
        and fold constant inputs values. When possible, run shape inference on the
        ONNX graph to determine tensor shapes.
        """
        for _ in range(3):
            count_before = len(self.graph.nodes)

            self.graph.cleanup().toposort()
            try:
                for node in self.graph.nodes:
                    for o in node.outputs:
                        o.shape = None
                model = gs.export_onnx(self.graph)
                model = shape_inference.infer_shapes(model)
                self.graph = gs.import_onnx(model)
            except Exception as e:
                self.logger.info(f"{self.prefix} Shape inference could not be performed at this time:\n{e}")
            try:
                self.graph.fold_constants(fold_shapes=True)
            except TypeError as ex:
                a = f"{self.prefix} This version of ONNX GraphSurgeon does not support folding shapes, " \
                      f"please upgrade your onnx_graphsurgeon module. Error:\n{ex}"
                self.logger.info(a)

            count_after = len(self.graph.nodes)
            if count_before == count_after:
                # No new folding occurred in this iteration, so we can stop for now.
                break

    def save(self, output_path):
        """
        Save the ONNX model to the given location.
        Args:
            output_path: Path pointing to the location where to write
                out the updated ONNX model.
        """
        self.graph.cleanup().toposort()
        model = gs.export_onnx(self.graph)
        onnx.save(model, output_path)
        self.logger.info(f"{self.prefix} saved ONNX model to {output_path}")

    def register_nms(self, *, score_thresh: float = 0.75, nms_thresh: float = 0.5, detections_per_img: int = 100,):
        """
        Register the ``EfficientNMS_TRT`` plugin node.
        NMS expects these shapes for its input tensors:
            - box_net: [batch_size, number_boxes, 4]
            - class_net: [batch_size, number_boxes, number_labels]
        Args:
            score_thresh (float): The scalar threshold for score (low scoring boxes are removed).
            nms_thresh (float): The scalar threshold for IOU (new boxes that have high IOU
                overlap with previously selected boxes are removed).
            detections_per_img (int): Number of best detections to keep after NMS.
        """

        self.infer()
        # Find the concat node at the end of the network
        op_inputs = self.graph.outputs
        op = "EfficientNMS_TRT"
        attrs = {
            "plugin_version": "1",
            "background_class": -1,  # no background class
            "max_output_boxes": detections_per_img,
            "score_threshold": score_thresh,
            "iou_threshold": nms_thresh,
            "score_activation": False,
            "box_coding": 0,
        }

        if self.precision == "fp32":
            dtype_output = np.float32
        elif self.precision == "fp16":
            dtype_output = np.float16
        else:
            raise NotImplementedError(f"Currently not supports precision: {self.precision}")

        # NMS Outputs
        output_num_detections = gs.Variable(name="num_dets",dtype=np.int32, shape=[self.batch_size, 1],)  # A scalar indicating the number of valid detections per batch image.
        output_boxes = gs.Variable(name="det_boxes",dtype=dtype_output, shape=[self.batch_size, detections_per_img, 4],)
        output_scores = gs.Variable(name="det_scores",dtype=dtype_output, shape=[self.batch_size, detections_per_img],)
        output_labels = gs.Variable(name="det_classes",dtype=np.int32, shape=[self.batch_size, detections_per_img],)
        op_outputs = [output_num_detections, output_boxes, output_scores, output_labels]

        # Create the NMS Plugin node with the selected inputs. The outputs of the node will also
        # become the final outputs of the graph.
        self.graph.layer(op=op, name="batched_nms", inputs=op_inputs, outputs=op_outputs, attrs=attrs)
        self.logger.info(f"{self.prefix} created NMS plugin '{op}' with attributes: \n{attrs}")

        self.graph.outputs = op_outputs
        self.infer()

    def save(self, output_path, onnx_MetaData=None):
        """
        Save the ONNX model to the given location.
        Args:
                @param output_path:
                @param onnx_MetaData:
        """
        self.graph.cleanup().toposort()
        model = gs.export_onnx(self.graph)
        if onnx_MetaData:
            for index, key in enumerate(onnx_MetaData):
                metadata = model.metadata_props.add()
                metadata.key = key
                metadata.value = str(onnx_MetaData[key])
            onnxmltools.utils.save_model(model, output_path)
        else:
            onnx.save(model, output_path)
        self.logger.info(f"{self.prefix} saved ONNX model to {output_path}")
